Declare an ONNX inference backend for basic-pitch - #18

Closes #12

Problem

The image builds fine but the container never starts on linux/arm64. Shimmy waits 15 s for the worker, the worker exits 1, and the container dies:

NameError: name '_default_model_type' is not defined
  File "/app/.venv/lib/python3.11/site-packages/basic_pitch/__init__.py", line 95

Basic Pitch ships no inference code of its own. It picks a backend at import from whichever of CoreML, TensorFlow, TFLite or ONNX Runtime it finds, and fails while defining ICASSP_2022_MODEL_PATH if it finds none. Those backends sit behind platform markers, so what you get depends on the machine doing the install:

On arm64 the tensorflow wheel is an empty shim that requires tensorflow-cpu-aws, which Poetry never resolved into our lock. CI and the AWS build are amd64, so neither ever sees this.

Approach

Tests first, in evaluation_function/backend_test.py. Before the change, on this macOS host:

test_onnx_runtime_is_installed              FAILED
test_basic_pitch_resolves_a_model_path      PASSED
test_model_loads_through_our_own_loader     PASSED
1 failed, 2 passed

Being explicit about what that shows: only the first test can fail on a developer machine, because macOS resolves coremltools and satisfies the other two. The container is the real reproduction, so I built the image from this branch before and after.

Before (import evaluation_function.evaluation inside the image):

NameError: name '_default_model_type' is not defined

After:

import OK in 1.2s

These tests live in their own file rather than in evaluation_test.py because they do not exercise this repository's logic. They assert that a usable inference backend is installed, which is a packaging concern rather than pipeline behaviour. The same reasoning applies to transport_test.py in #17; if both land, folding the two into a single environment_test.py would be reasonable.

Change

One dependency line in pyproject.toml, plus the relock:

basic-pitch = { version = "^0.4.0", extras = ["onnx"] }

The lock change is confined to onnxruntime and its flatbuffers dependency.

Verification

Correction to the issue

I wrote in #12 that this would also shrink the image. That was wrong. The onnx extra adds a backend but does not remove tensorflow, which basic-pitch depends on unconditionally for non-Darwin platforms. The arm64 image grew from 1.4 GB to 1.46 GB.

Dropping TensorFlow in favour of ONNX Runtime alone would genuinely shrink the amd64 image, but that changes which backend runs in production and needs its own transcription-accuracy check, so it is out of scope here. Worth a separate issue if the image size matters.

🤖 Generated with Claude Code